import os
import json
import pickle
from typing import List
from pathlib import Path
from tqdm import tqdm

from dotenv import load_dotenv
from langchain.embeddings import OpenAIEmbeddings
from langchain.vectorstores import FAISS
from langchain.text_splitter import CharacterTextSplitter
from langchain.retrievers import BM25Retriever
from langchain_community.embeddings import HuggingFaceEmbeddings, DashScopeEmbeddings
from langchain_core.documents import Document

from langchain.schema import Document

class BM25Ingestor:
    def __init__(self):
        pass

    def process_reports(self, all_reports_dir: Path, output_dir: Path):
        """
        加载 JSON 报告，提取文本与 metadata，构建 BM25Retriever 并保存。

        Args:
            all_reports_dir (Path): JSON 报告目录
            output_index_path (Path): 保存的 bm25.pkl 文件路径
        """
        all_report_paths = list(all_reports_dir.glob("*.json"))
        all_documents: List[Document] = []
        output_dir.mkdir(parents=True, exist_ok=True)


        for report_path in tqdm(all_report_paths, desc="Processing reports"):
            with open(report_path, 'r', encoding='utf-8') as f:
                report_data = json.load(f)

            text_chunks = [chunk["text"] for chunk in report_data["content"]["chunks"]]
            metadata = report_data["metainfo"]

            for text in text_chunks:
                all_documents.append(Document(
                    page_content=text,
                    metadata=metadata.copy()  # 每段文本独立 metadata
                ))

        print(f"构建 BM25Retriever，文档数：{len(all_documents)}")
        retriever = BM25Retriever.from_documents(all_documents)
        bm25_file_path = output_dir / "bm25.pkl"
        with open(bm25_file_path, 'wb') as f:
            pickle.dump(retriever, f)

        print(f"BM25 索引保存至：{output_dir}")


class VectorDBIngestor:
    def __init__(self, checkpoint="./jinaai/jina-embeddings-v2-base-zh"):
        """
        初始化 VectorDBIngestor 类的实例。
        加载环境变量并初始化 OpenAI 嵌入模型。
        """
        # 加载 .env 文件中的环境变量
        load_dotenv()

        model_name = checkpoint
        model_kwargs = {"device": "cpu"}
        encode_kwargs = {"normalize_embeddings": False}
        embeddings = HuggingFaceEmbeddings(
            model_name=model_name,
            model_kwargs=model_kwargs,
            encode_kwargs=encode_kwargs,
        )

        self.embeddings = embeddings
        # 初始化 OpenAI 嵌入模型
        # self.embeddings = OpenAIEmbeddings(
        #     # 从环境变量中获取 OpenAI API 密钥
        #     api_key=os.getenv("OPENAI_API_KEY"),
        #     # 指定使用的 OpenAI 嵌入模型
        #     model="text-embedding-3-large"
        # )

    def _process_report(self, report: dict):
        """
        处理单个报告，将报告中的文本块转换为向量并构建向量数据库。

        Args:
            report (dict): 包含报告内容的字典。

        Returns:
            FAISS: 基于报告文本块创建的 FAISS 向量数据库实例。
        """
        # 从报告数据中提取文本块
        text_chunks = [chunk['text'] for chunk in report['content']['chunks']]
        # 使用提取的文本块和 OpenAI 嵌入模型创建 FAISS 向量数据库
        vector_db = FAISS.from_texts(text_chunks, self.embeddings)
        return vector_db

    # def process_reports(self, all_reports_dir: Path, output_dir: Path):
    #     """
    #     处理所有报告文件并保存单独的向量数据库。

    #     Args:
    #         all_reports_dir (Path): 包含 JSON 报告文件的目录路径。
    #         output_dir (Path): 用于保存向量数据库的目录路径。
    #     """
    #     # 获取指定目录下所有 JSON 文件的路径列表
    #     all_report_paths = list(all_reports_dir.glob("*.json"))
    #     # 创建输出目录，如果目录已存在则不会报错
    #     output_dir.mkdir(parents=True, exist_ok=True)

    #     # 遍历所有报告文件路径，使用 tqdm 显示进度条
    #     for report_path in tqdm(all_report_paths, desc="Processing reports"):
    #         # 加载报告文件
    #         with open(report_path, 'r', encoding='utf-8') as file:
    #             report_data = json.load(file)
    #         # 处理单个报告并创建向量数据库
    #         vector_db = self._process_report(report_data)
    #         # 获取报告的 sha1 名称
    #         sha1_name = report_data["metainfo"]["sha1_name"]
    #         # 生成保存向量数据库的文件路径
    #         faiss_file_path = output_dir / f"{sha1_name}.faiss"
    #         # 将向量数据库保存到本地文件
    #         vector_db.save_local(faiss_file_path)

    #     # 打印处理的报告文件数量
    #     print(f"Processed {len(all_report_paths)} reports")

    def process_reports_all(self, all_reports_dir: Path, output_dir: Path):
        """
        处理所有报告文件并将所有数据存储到一个向量数据库中，同时保存元数据。

        Args:
            all_reports_dir (Path): 包含 JSON 报告文件的目录路径。
            output_dir (Path): 用于保存向量数据库的目录路径。
        """
        # 获取指定目录下所有 JSON 文件的路径列表
        all_report_paths = list(all_reports_dir.glob("*.json"))
        # 创建输出目录，如果目录已存在则不会报错
        output_dir.mkdir(parents=True, exist_ok=True)

        all_text_chunks = []
        all_metadata = []
        all_docs = []
        # 遍历所有报告文件路径，使用 tqdm 显示进度条
        for report_path in tqdm(all_report_paths, desc="Processing reports"):
            # 加载报告文件
            with open(report_path, 'r', encoding='utf-8') as file:
                report_data = json.load(file)
            # 从报告数据中提取文本块
            text_chunks = [chunk['text'] for chunk in report_data['content']['chunks']]
            all_text_chunks.extend(text_chunks)
            
            # 提取元数据，这里假设每个文本块共享相同的元数据
            metainfo = report_data["metainfo"]
            for text in text_chunks:
                docs = Document(
                    page_content=text,
                    metadata=metainfo.copy()  # 每段文本独立 metadata
                )
                all_docs.append(docs)
                all_metadata.append(metainfo.copy())



        # 处理所有文本块并创建一个向量数据库，同时传入元数据
        # vector_db = FAISS.from_texts(all_text_chunks, self.embeddings, metadatas=all_metadata)
        vector_db = FAISS.from_documents(all_docs, self.embeddings)

        # 生成保存向量数据库的文件路径
        faiss_file_path = output_dir / "reports_faiss"
        # 将向量数据库保存到本地文件
        vector_db.save_local(faiss_file_path)

        # 打印处理的报告文件数量
        print(f"Processed {len(all_report_paths)} reports and saved to a single vector database with metadata.")